package com.iraytek.rtsplib;

import android.media.MediaCodec;
import android.util.Log;

import org.jetbrains.annotations.NotNull;

import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

import pedro.rtsp.rtcp.BaseSenderReport;
import pedro.rtsp.rtp.packets.AacPacket;
import pedro.rtsp.rtp.packets.AudioPacketCallback;
import pedro.rtsp.rtp.packets.BasePacket;
import pedro.rtsp.rtp.packets.H264Packet;
import pedro.rtsp.rtp.packets.H265Packet;
import pedro.rtsp.rtp.packets.VideoPacketCallback;
import pedro.rtsp.rtp.sockets.BaseRtpSocket;
import pedro.rtsp.rtsp.Protocol;
import pedro.rtsp.rtsp.RtpFrame;
import pedro.rtsp.utils.BitrateManager;
import pedro.rtsp.utils.ConnectCheckerRtsp;
import pedro.rtsp.utils.RtpConstants;

/**
 * created by yhy on 2021/1/12
 * Used to send RTSP packet, reference {@link pedro.rtsp.rtsp.RtspSender}
 */
public class RtspSender implements VideoPacketCallback, AudioPacketCallback {
    private static final String TAG = "RtspSender";
    private boolean running;
    private boolean isEnableLogs = true;
    private long audioFramesSent;
    private long videoFramesSent;
    private long droppedAudioFrames;
    private long droppedVideoFrames;
    private int defaultCacheSize = 10 * 1024 * 1024 / RtpConstants.MTU;

    private ConnectCheckerRtsp connectCheckerRtsp;
    private BaseRtpSocket rtpSocket;
    private BaseSenderReport baseSenderReport;
    private BasePacket videoPacket;
    private AacPacket aacPacket;
    private BlockingQueue<RtpFrame> rtpFrameBlockingQueue;
    private BitrateManager bitrateManager;
    private Thread thread;

    public RtspSender(ConnectCheckerRtsp connectCheckerRtsp) {
        this.connectCheckerRtsp = connectCheckerRtsp;
        bitrateManager = new BitrateManager(connectCheckerRtsp);
        rtpFrameBlockingQueue = new LinkedBlockingQueue<>(defaultCacheSize);
    }


    @Override
    public void onAudioFrameCreated(@NotNull RtpFrame rtpFrame) {
        try {
            rtpFrameBlockingQueue.add(rtpFrame);
        } catch (IllegalStateException e) {
            e.printStackTrace();
            droppedAudioFrames++;
        }
    }

    @Override
    public void onVideoFrameCreated(@NotNull RtpFrame rtpFrame) {
        try {
            rtpFrameBlockingQueue.add(rtpFrame);
        } catch (IllegalStateException e) {
            e.printStackTrace();
            droppedVideoFrames++;
        }
    }

    public void setSocketsInfo(Protocol protocol, int[] videoPorts, int[] audioPorts) {
        rtpSocket = BaseRtpSocket.getInstance(protocol, videoPorts[0], audioPorts[0]);
        baseSenderReport = BaseSenderReport.getInstance(protocol, videoPorts[1], audioPorts[1]);
    }

    public void setVideoInfo(byte[] sps, byte[] pps, byte[] vps) {
        videoPacket = vps == null ? new H264Packet(sps, pps, this) :
                new H265Packet(sps, pps, vps, this);
    }

    public void setAudioInfo(int sampleRate) {
        aacPacket = new AacPacket(sampleRate, this);
    }

    public void setDataStream(OutputStream outputStream, String host) {
        if (rtpSocket != null) {
            rtpSocket.setDataStream(outputStream, host);
        }
        if (baseSenderReport != null) {
            baseSenderReport.setDataStream(outputStream, host);
        }
    }

    public void setVideoPorts(int rtpPort, int rtcpPort) {
        if (videoPacket != null) {
            videoPacket.setPorts(rtpPort, rtcpPort);
        }
    }

    public void setAudioPorts(int rtpPort, int rtcpPort) {
        if (aacPacket != null) {
            aacPacket.setPorts(rtpPort, rtcpPort);
        }
    }

    public void sendVideoFrame(@NotNull ByteBuffer h264Buffer, @NotNull MediaCodec.BufferInfo info) {
        if (running && videoPacket != null) videoPacket.createAndSendPacket(h264Buffer, info);
    }

    public void sendAudioFrame(@NotNull ByteBuffer aacBuffer, @NotNull MediaCodec.BufferInfo info) {
        if (running && aacPacket != null) aacPacket.createAndSendPacket(aacBuffer, info);
    }

    public void start() {
        thread = new Thread(() -> {
            running = true;
            while (!thread.isInterrupted()) {
                try {
                    RtpFrame rtpFrame = rtpFrameBlockingQueue.poll(1, TimeUnit.SECONDS);
                    if (rtpFrame == null) {
                        Log.w(TAG, "Skipping iteration, frame null");
                        continue;
                    }
                    sendFrame(rtpFrame, isEnableLogs);
                    //bytes to bits
                    bitrateManager.calculateBitrate(rtpFrame.getLength() * 8);
                    if (rtpFrame.isVideoFrame()) {
                        videoFramesSent++;
                    } else {
                        audioFramesSent++;
                    }
                    if (baseSenderReport != null) {
                        baseSenderReport.update(rtpFrame, isEnableLogs);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                    if (e instanceof InterruptedException) {
                        connectCheckerRtsp.onConnectionFailedRtsp("Error send packet, " + e.getMessage());
                    }
                    Log.e(TAG, "send error: ", e);
                    return;
                }
            }
        });
        thread.start();
    }

    /**
     * Send frame to network by socket, override it if you want to send frame by other methods.
     *
     * @param rtpFrame
     * @param isEnableLogs
     * @throws IOException
     */
    protected void sendFrame(RtpFrame rtpFrame, boolean isEnableLogs) throws IOException {
        if (rtpSocket != null) {
            rtpSocket.sendFrame(rtpFrame, isEnableLogs);
        }
    }

    public void stop() {
        running = false;
        if (thread != null) {
            thread.interrupt();
            try {
                thread.join(100);
            } catch (Exception e) {
            }
            thread = null;
        }

        rtpFrameBlockingQueue.clear();
        if (baseSenderReport != null) {
            baseSenderReport.reset();
            baseSenderReport.close();
        }
        if (rtpSocket != null) {
            rtpSocket.close();
        }
        if (aacPacket != null) {
            aacPacket.reset();
        }
        if (videoPacket != null) {
            videoPacket.reset();
        }
        resetSentAudioFrames();
        resetSentVideoFrames();
        resetDroppedAudioFrames();
        resetDroppedVideoFrames();
    }

    public void resetDroppedVideoFrames() {
        droppedVideoFrames = 0;
    }

    public void resetDroppedAudioFrames() {
        droppedAudioFrames = 0;
    }

    public void resetSentVideoFrames() {
        videoFramesSent = 0;
    }

    public void resetSentAudioFrames() {
        audioFramesSent = 0;
    }

    public long getDroppedAudioFrames() {
        return droppedAudioFrames;
    }

    public long getDroppedVideoFrames() {
        return droppedVideoFrames;
    }

    public int getCacheSize() {
        return rtpFrameBlockingQueue.size();
    }

    public long getSentAudioFrames() {
        return audioFramesSent;
    }

    public long getSentVideoFrames() {
        return audioFramesSent;
    }


    public boolean hasCongestion() {
        float size = rtpFrameBlockingQueue.size();
        float remaining = rtpFrameBlockingQueue.remainingCapacity();
        float capacity = size + remaining;
        return size >= capacity * 0.2f; //more than 20% queue used. You could have congestion
    }

    public void resizeCache(int newSize) {
        if (newSize < rtpFrameBlockingQueue.size() - rtpFrameBlockingQueue.remainingCapacity()) {
            throw new RuntimeException("Can't fit current cache inside new cache size");
        }
        BlockingQueue<RtpFrame> tempQueue = new LinkedBlockingQueue(newSize);
        rtpFrameBlockingQueue.drainTo(tempQueue);
        rtpFrameBlockingQueue = tempQueue;
    }

    public void setLogs(boolean enable) {
        isEnableLogs = enable;
    }
}
